import numpy as np
import cv2
import pyopencl as cl
import time
import matplotlib.pyplot as plt
from skimage import color
from PIL import Image

# 定义终端调色板 (不含灰色)
TERMINAL_PALETTE_RGB = np.array(
    [
        [0, 0, 0],  # 黑
        [128, 0, 0],  # 暗红
        [0, 128, 0],  # 暗绿
        [128, 128, 0],  # 暗黄
        [0, 0, 128],  # 暗蓝
        [128, 0, 128],  # 暗紫
        [0, 128, 128],  # 暗青
        [255, 100, 100],  # 亮粉红
        [100, 255, 100],  # 亮青绿
        [255, 255, 100],  # 亮黄
        [100, 100, 255],  # 亮蓝
        [255, 100, 255],  # 亮紫
        [100, 255, 255],  # 亮青
        [255, 200, 150],  # 肉色
        [200, 255, 200],  # 薄荷绿
        [255, 255, 255],  # 白
    ],
    dtype=np.uint8,
)

# 将RGB调色板转换为LAB颜色空间
TERMINAL_PALETTE_LAB = color.rgb2lab(TERMINAL_PALETTE_RGB.reshape(1, -1, 3)).reshape(
    -1, 3
)


def enhance_brightness(image, brightness_factor=1.3, contrast_factor=1.2):
    """
    增强图像亮度和对比度
    :param image: 输入图像 (RGB格式)
    :param brightness_factor: 亮度增强因子 (1.0为原始亮度)
    :param contrast_factor: 对比度增强因子 (1.0为原始对比度)
    :return: 增强后的图像
    """
    # 转换为浮点型以进行精确计算
    img_float = image.astype(np.float32) / 255.0

    # 应用亮度调整
    img_bright = np.clip(img_float * brightness_factor, 0, 1)

    # 应用对比度调整
    mean_val = np.mean(img_bright, axis=(0, 1))
    img_contrast = np.clip((img_bright - mean_val) * contrast_factor + mean_val, 0, 1)

    # 转换回8位整数
    result = (img_contrast * 255).astype(np.uint8)

    return result


def apply_terminal_filter(image_path, brightness=1.3, contrast=1.2):
    # 读取图像并转换为RGB
    # img = cv2.imread(image_path)
    img = Image.open(image_path)
    # img.convert("LAB")
    img = np.array(img)
    if img is None:
        raise FileNotFoundError(f"无法加载图像: {image_path}")
    # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    # 增强亮度和对比度
    img_enhanced = enhance_brightness(img, brightness, contrast)

    # 将图像转换为LAB颜色空间
    img_lab = color.rgb2lab(img_enhanced)

    height, width, _ = img_enhanced.shape

    # 设置OpenCL环境
    ctx = cl.create_some_context()
    queue = cl.CommandQueue(ctx)

    # 创建OpenCL程序 (使用LAB颜色空间)
    kernel_code = open("./kernel.cl", "r", encoding="utf-8").read()
    """"""

    # 编译OpenCL程序
    program = cl.Program(ctx, kernel_code).build()

    # 准备输入输出缓冲区
    # 将LAB图像展平为一维数组
    img_lab_flat = img_lab.astype(np.float32).ravel()
    input_buf = cl.Buffer(
        ctx, cl.mem_flags.READ_ONLY | cl.mem_flags.COPY_HOST_PTR, hostbuf=img_lab_flat
    )
    output_buf = cl.Buffer(ctx, cl.mem_flags.WRITE_ONLY, img_enhanced.nbytes)

    # 执行GPU内核
    start_time = time.time()
    program.terminal_filter(
        queue,
        (width, height),
        None,
        input_buf,
        output_buf,
        np.int32(width),
        np.int32(height),
    )
    queue.finish()
    gpu_time = time.time() - start_time

    # 获取结果
    result = np.empty_like(img_enhanced)
    cl.enqueue_copy(queue, result, output_buf)

    # 转换为BGR并保存
    result_bgr = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
    cv2.imwrite("laboutput.png", result_bgr)

    # 显示结果对比
    # plt.figure(figsize=(18, 6))
    # plt.subplot(131)
    # plt.title("原始图像")
    # plt.imshow(img)
    # plt.axis("off")

    # plt.subplot(132)
    # plt.title("增强后图像")
    # plt.imshow(img_enhanced)
    # plt.axis("off")

    # plt.subplot(133)
    # plt.title("终端滤镜效果 (LAB空间)")
    # plt.imshow(result)
    # plt.axis("off")
    # plt.savefig("result_comparison.png", bbox_inches="tight")
    # plt.close()

    # 打印性能信息
    print(f"GPU处理时间: {gpu_time:.4f}秒")
    print(f"图像尺寸: {width}x{height} 像素")
    print(f"输出已保存为 output.png")
    # print(f"结果对比图保存为 result_comparison.png")

    return result_bgr


if __name__ == "__main__":
    import sys

    if len(sys.argv) != 2:
        # print("使用方法: python terminal_filter_lab.py <图片路径>")
        # sys.exit(1)
        input_image = "img.jpeg"
    else:
        input_image = sys.argv[1].strip()

    # 可调节参数：亮度和对比度增强因子
    # 对于暗图像，可以增加这些值 (如 brightness=1.5, contrast=1.3)
    result = apply_terminal_filter(input_image, brightness=1.2, contrast=4)
